import argparse


#! QP类型
# enum ibv_qp_type:
#         IBV_QPT_RC
#         IBV_QPT_UC
#         IBV_QPT_UD
#         IBV_QPT_RAW_PACKET
#         IBV_QPT_XRC_SEND
#         IBV_QPT_XRC_RECV
#         IBV_QPT_DRIVER
from pyverbs.enums import IBV_QPT_UD # type: ignore
from pyverbs.enums import IBV_QPT_RC # type: ignore


#! RDMA操作类型
# enum ibv_wr_opcode:
#         IBV_WR_RDMA_WRITE
#         IBV_WR_RDMA_WRITE_WITH_IMM
#         IBV_WR_SEND
#         IBV_WR_SEND_WITH_IMM
#         IBV_WR_RDMA_READ
#         IBV_WR_ATOMIC_CMP_AND_SWP
#         IBV_WR_ATOMIC_FETCH_AND_ADD
#         IBV_WR_LOCAL_INV
#         IBV_WR_BIND_MW
#         IBV_WR_SEND_WITH_INV
#         IBV_WR_TSO
#         IBV_WR_FLUSH
#         IBV_WR_ATOMIC_WRITE
from pyverbs.enums import IBV_WR_SEND        # type: ignore
from pyverbs.enums import IBV_WR_RDMA_WRITE  # type: ignore
from pyverbs.enums import IBV_WR_RDMA_READ   # type: ignore



#! RDMA QP Type(连接类型)
# RC: 可靠连接
# UD: 不可靠数据报
rdma_qpt  = {'RC': IBV_QPT_RC, 'UD': IBV_QPT_UD}


#! RDMA 操作符
rdma_ops = {
            'SEND' : IBV_WR_SEND,        # 双端操作Send
            'READ' : IBV_WR_RDMA_READ,   # 单端操作RDMA_READ
            'WRITE': IBV_WR_RDMA_WRITE,  # 单端操作RDMA_WRITE
        }


class DictAction(argparse.Action):
    def __call__(self, parsers, namespace, values, option_string=None):
        setattr(namespace, self.dest, self.choices.get(values))


class ArgsParser(object):
    def __init__(self):
        self.args = None

    @staticmethod
    def parse_args():
        arg_parser = argparse.ArgumentParser(description="A Python Test Program for RDMA.")
        arg_parser.add_argument('server_ip', type=str, nargs='?',
                                help='Server端的IP地址(Client only).')
        
        #! RDMA在正式数据传输前，需要进行建连，以交换元数据, 如内存地址、RKey等(); 
        #  RDMA支持两种方式建连:
        #   - 使用传统的以太网套接字;
        #   - 使用RDMA自带的Connection Management协议(简称CM);
        arg_parser.add_argument('-C', '--use_cm', action='store_true',
                                help='是否使用CM协议进行握手.')
        
        #! 通信使用的RDMA HCA(Host Channel Adapter)设备名称, 即所谓的RDM网卡
        # 可通过命令查询: (base) [yangxianpku@localhost rdma]$ ibv_devinfo
        arg_parser.add_argument('-d', '--ib_dev', required=True,
                                help='RDMA device to run the tests on.')
        
        #  - SGE表示一个数据段(Data Segment)， 其本质是一个数据结构
        #       struct ibv_sge {
        #           uint64_t addr;   // 数据段的起始地址
        #           uint32_t length; // 数据段的长度
        #           uint32_t lkey;   // 本地密钥
        #       };
        #
        # 
        # - SGL（Scatter/Gather List）
        #   SGL 是 Scatter/Gather List 的缩写，是一个由多个 SGE 组成的数组。SGL 的主要作用是：
        #   - 分散（Scatter）：允许一个连续的缓冲区（Buffer）分散到多个不连续的内存区域。
        #   - 聚合（Gather）：允许多个不连续的内存区域聚合到一个连续的缓冲区。
        # 
        #   在 RDMA 编程中，SGL 是基本的数据组织形式。例如，在使用 ibv_post_send 或 ibv_post_recv 
        #   等函数时，需要提供一个 sg_list 数组，该数组即为 SGL 的核心部分


        #  3. SGE 与 SGL 的关系
        #    - SGE 是 SGL 的基本组成单元：SGL 是一个数组，其中的每个元素都是一个 SGE。
        #    - SGL 提供灵活性：通过 SGL，RDMA 硬件可以在内存中处理分散的数据段，并通过聚合操作将它们传
        #      输为连续的数据流。
        #! Satter/Gather的深度, 即SGL(Satter/Gather List)中SGE(Satter/Gather Element)的数量
        arg_parser.add_argument('-G', '--sg_depth', type=int, default=1,
                                help='SGL中SGE的数量.')
        

        arg_parser.add_argument('-I', '--inline_size', type=int, default=0,
                                help='内联发送的最大消息大小.')
        

        #! 设置RDMA的Path MTU
        # MTU: Maximum Transfer Unit, 即传输过程中单个数据包的大小(数据主题，不含包头的协议部分)
        # 更大的MTU可以减少包的数量, 从而提高带宽， 但也更容易丢包, 因此MTU的设置需要根据网络环境进行调整
        arg_parser.add_argument('-m', '--mtu', type=int, default=4, choices=range(0, 5),
                                help='The Path MTU 0 ~ 4 (256/512/1024/2048/4096).')
        

        arg_parser.add_argument('-n', '--iters', type=int, default=1,
                                help='Number of exchanges (at least 5, default 1).')
        

        # 服务端监听或者客户端连接的端口
        arg_parser.add_argument('-p', '--port', type=int, default=18515,
                                help='服务端监听或者客户端连接的端口(default 18515).')
        
        # 发送队列的大小
        arg_parser.add_argument('-t', '--tx_depth', type=int, default=16,
                                help='Size of tx queue (default 16).')
        

        # 接收队列的大小
        # 如果使用共享接收队列， 则代指共享接收队列的大小(Shared RQ Size)
        arg_parser.add_argument('-r', '--rx_depth', type=int, default=16,
                                help='Size of rx queue (default 16). '
                                'If using srq, rx-depth controls max-wr '
                                'size of the srq.')
        

        # MR区域的大小
        arg_parser.add_argument('-s', '--mr_size', type=int, default=67108864,
                                help='交换的消息大小(默认64MB).')
        

        arg_parser.add_argument('-S', '--sl', type=int, default=0,
                                help='服务等级 (default 0).')
        

        # RDMA操作类型
        arg_parser.add_argument('-o', '--op_type', 
                                default=IBV_WR_SEND, 
                                choices=rdma_ops, 
                                action=DictAction,
                                help='RDMA操作类型.'
                            )
        
        

        # QP的类型
        arg_parser.add_argument('-T', '--qp_type', 
                                type=str, 
                                default=IBV_QPT_RC, 
                                choices=rdma_qpt, 
                                action=DictAction,
                                help='QP类型.'
                            )
        

        arg_parser.add_argument('-v', '--version', 
                                action='version', 
                                version='%(prog)s version : v0.01',
                                help='Show the version')
        

        # GID: GID是HCA厂商赋予HCA设备的全局唯一标识符;
        # LID: LID是子网管理器分配给HCA端口的唯一地址，可能随网络或重启而变化，用于IB网络,RoCE网络无效; 
        # 可通过命令查询:
        # (base) [yangxianpku@localhost rdma]$ ibv_devinfo
        # hca_id: rocep0s6
        #         transport:                      InfiniBand (0)
        #         fw_ver:                         32.42.1000
        #         node_guid:                      9c63:c003:005c:517c
        #         sys_image_guid:                 9c63:c003:005c:517c
        #         vendor_id:                      0x02c9
        #         vendor_part_id:                 41692
        #         hw_ver:                         0x1
        #         board_id:                       MT_0000001070
        #         phys_port_cnt:                  1
        #                 port:   1
        #                         state:                  PORT_ACTIVE (4)
        #                         max_mtu:                4096 (5)
        #                         active_mtu:             1024 (3)
        #                         sm_lid:                 0
        #                         port_lid:               0
        #                         port_lmc:               0x00
        #                         link_layer:             Ethernet (1)
        arg_parser.add_argument('-x', '--gid_index', type=int, default=1,
                                help='与远程端点通信的源GID的索引号')

        # TODO: Check the range and raise an error, for example:
        # https://stackoverflow.com/questions/18700634/python-argparse-integer-condition-12

        return vars(arg_parser.parse_args())


parser = ArgsParser()

